/**
 * Java implementation of a Red-Black Tree of type T.
 * 
 * @author Jonathan E. Landrum <me@jonlandrum.com>
 * @since 2016-01-22
 */

package com.jonlandrum.selfbalancingtree.redblacktree;

import com.jonlandrum.selfbalancingtree.Node;
import com.jonlandrum.selfbalancingtree.SelfBalancingTree;

public class RedBlackTree<T extends Comparable<T>> extends SelfBalancingTree<T> {
    /*
     * Constructors
     */
    public RedBlackTree() {
        super();
    }
    
    public RedBlackTree(T d) {
        super(d);
        this.getRoot().setLeftChild(new Node<T>());
        this.getRoot().setRightChild(new Node<T>());
    }
    
    public RedBlackTree(Node<T> n) {
        super(n);
        this.getRoot().setLeftChild(new Node<T>());
        this.getRoot().setRightChild(new Node<T>());
    }
    
    /*
     * Operations
     */
    @Override
    public Node<T> addElement(T d) {
        return this.addElement(new Node<T>(d));
    }
    
    @Override
    public Node<T> addElement(Node<T> n) {
        Node<T> result = super.addElement(n);
        if (result.isEmpty()) { return result; }
        result.swapColor();
        result.setLeftChild(new Node<T>());
        result.setRightChild(new Node<T>());
        Node<T> temp = result;
        while (temp != this.getRoot() && temp.getParent().getColor() == 0) {
            if (temp.getParent().getSibling().getColor() == 0) {
                temp.getParent().swapColor();
                temp.getParent().getSibling().swapColor();
                temp.getParent().getParent().swapColor();
                temp = temp.getParent();
            } else if (temp.getParent().isLeftChild()) {
                if (temp.isRightChild()) {
                    this.rotateLeft(temp.getParent());
                    temp = temp.getLeftChild();
                }
                this.rotateRight(temp.getParent().getParent());
                temp.getParent().swapColor();
                temp.getParent().getRightChild().swapColor();
            } else if (temp.getParent().isRightChild()) {
                if (temp.isLeftChild()) {
                    this.rotateRight(temp.getParent());
                    temp = temp.getRightChild();
                }
                this.rotateLeft(temp.getParent().getParent());
                temp.getParent().swapColor();
                temp.getParent().getLeftChild().swapColor();
            }
            temp = temp.getParent();
        }
        if (this.getRoot().getColor() == 0) {
            this.getRoot().swapColor();
        }
        return result;
    }
    
    @Override
    public Node<T> removeElement(T d) {
        return this.removeElement(new Node<T>(d));
    }
    
    @Override
    public Node<T> removeElement(Node<T> n) {
        Node<T> result = super.findElement(n);
        if (result.isEmpty()) {
            return result;
        } else if (result.getColor() == 0) {
            return super.removeElement(result);
        } else {
            Node<T> successor = new Node<T>();
            if (result.hasRightChild()) {
                successor = result.getRightChild();
                while (successor.hasLeftChild()) {
                    successor = successor.getLeftChild();
                }
                result.setData(successor);
                if (successor.hasRightChild()) {
                    if (successor.isLeftChild()) {
                        successor.getParent().setLeftChild(successor.getRightChild());
                    } else {
                        successor.getParent().setRightChild(successor.getRightChild());
                    }
                    successor.getRightChild().setParent(successor.getParent());
                } else {
                    if (successor.isLeftChild()) {
                        successor.getParent().setLeftChild(new Node<T>());
                    } else {
                        successor.getParent().setRightChild(new Node<T>());
                    }
                }
            } else if (result.hasLeftChild()) {
                successor = result.getLeftChild();
                while (successor.hasRightChild()) {
                    successor = successor.getRightChild();
                }
                result.setData(successor);
            } else {
                result = result.getParent();
            }
        }
        return result;
    }
}